import tensorflow as tf
from keras.datasets import mnist
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

# 加载MNIST数据集
(X_train, y_train), (X_test, y_test) = mnist.load_data()

# 将图像数据展平为一维数组并归一化
X_train = X_train.reshape(X_train.shape[0], -1) / 255.0
X_test = X_test.reshape(X_test.shape[0], -1) / 255.0

# 将标签转换为one-hot编码
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)
y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)

# 初始化变量以存储最佳准确率和相应的模型
best_accuracy = 0
best_model = None

# 初始化一个列表以存储每个epoch的准确率
accuracies = []

# 构建MLP模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 使用tqdm创建进度条,并遍历每个epoch
for epoch in tqdm(range(40), desc="训练进度", unit="epoch"):
    # 训练模型
    history = model.fit(X_train, y_train, batch_size=128, epochs=1, validation_data=(X_test, y_test), verbose=0)
    
    # 获取当前epoch的准确率
    accuracy = history.history['val_accuracy'][0]
    accuracies.append(accuracy)
    
    # 更新最佳准确率和相应的模型
    if accuracy > best_accuracy:
        best_accuracy = accuracy
        best_model = model
        
# 保存最佳模型
best_model.save('best_mlp_model.h5')

# 打印最佳准确率
print(f'最佳准确率: {best_accuracy}')

# 绘制准确率变化图并保存为pdf文件
plt.figure(figsize=(12, 6))
plt.plot(range(1, len(accuracies)+1), accuracies, color='blue', linestyle='dashed', marker='o', markerfacecolor='red', markersize=10)
plt.title('Accuracy Rate Epoch')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.savefig('accuracy_plot.pdf')